from __future__ import absolute_import
import os
import logging
import re
from obs import ObsClient
from conf.conf import (ObsAK, ObsEndpoint, ObsSK, BucketName, ObjectName, FileDir)
obs_client = ObsClient(access_key_id=ObsAK,
                       secret_access_key=ObsSK, server=ObsEndpoint)

logger = logging.getLogger(__name__)


def download_from_obs(bucket_name, object_key, file_dir):
    object_list = object_key.split(',')
    for objects in object_list:
        if objects == 'anonymize':
            make_anonymize(bucket_name, file_dir)
        elif objects == 'scene':
            make_scene(bucket_name, file_dir)
        elif objects == 'original':
            make_original(bucket_name, file_dir)
        else:
            logging.info("error object_key")


def make_anonymize(bucket_name, root_dir):
    resp = obs_client.listObjects(bucket_name)
    if resp.status < 300:
        for content in resp.body.contents:
            if not re.search(r'anonymize', content.key) or re.search(root_dir, content.key):
                continue
            key = content.key.rsplit('/', 1)
            file_dir = key[0]
            if not root_dir == "":
                file_dir = '{}/{}'.format(root_dir, file_dir)
            file_name = key[1]
            if file_name == '':
                if not os.path.exists(file_dir):
                    os.makedirs(file_dir)
            else:
                download_path = '{}/{}'.format(file_dir, file_name)
                download_file(bucket_name, content.key, download_path, content.storageClass)
                return
    else:
        logger.error("查询obs数据失败%s", resp.errorCode)


def make_scene(bucket_name, root_dir):
    resp = obs_client.listObjects(bucket_name)
    download_folder = ""
    if resp.status < 300:
        for content in resp.body.contents:
            if not re.search(r'scene', content.key) or not re.search(root_dir, content.key):
                continue
            if content.storageClass == "COLD":
                continue
            if not content.key.endswith('/'):
                key = content.key.rsplit('/', 1)
                file_name = key[1]
                file_dir = key[0]
                # 保证下载文件同属于一个文件夹，即完整的场景数据包
                if download_folder not in ('', file_dir):
                    return
                cur_dir = file_dir
                if not os.path.exists(cur_dir):
                    os.makedirs(cur_dir)
                download_path = '{}/{}'.format(cur_dir, file_name)
                download_folder = cur_dir
                download_file(bucket_name, content.key, download_path, content.storageClass)

    else:
        logger.error("查询obs数据失败%s", resp.errorCode)


def make_original(bucket_name, root_dir):
    resp = obs_client.listObjects(bucket_name)
    if resp.status < 300:
        for content in resp.body.contents:
            if not re.search(r'original', content.key) or re.search(root_dir, content.key):
                continue
            key = content.key.rsplit('/', 1)
            file_dir = key[0]
            if not root_dir == "":
                file_dir = '{}/{}'.format(root_dir, file_dir)
            file_name = key[1]
            if file_name == '':
                if not os.path.exists(file_dir):
                    os.makedirs(file_dir)
            else:
                download_path = '{}/{}'.format(file_dir, file_name)
                download_file(bucket_name, content.key, download_path, content.storageClass)
                return
    else:
        logger.error("查询obs数据失败%s", resp.errorCode)


def download_file(bucket_name, object_key, download_path, storage_class):
    if storage_class == "COLD":
        res = check_cold_file(bucket_name, object_key)
        if res:
            resp = obs_client.getObject(bucket_name, object_key, download_path)
        else:
            return 0
    else:
        resp = obs_client.getObject(bucket_name, object_key, download_path)
    if resp.status < 300:
        return 1
    else:
        logger.error("下载obs数据失败%s", resp.errorCode)
        return 0


def check_cold_file(bucket_name, object_key):
    resp = obs_client.getObjectMetadata(bucket_name, object_key)
    if resp.status < 300:
        if resp.body.restore is None:
            return 0
        elif resp.body.restore.find("false") != -1:  # 已恢复才能下载
            return 1
        else:
            return 0
    else:
        logger.error("查询obs数据失败%s", resp.errorCode)
        return 0


if __name__ == '__main__':
    download_from_obs(BucketName, ObjectName, FileDir)  # 下载obs文件至指定文件夹 会下载以FileDir为根目录的OBS文件
